import os
import re
from typing import Dict, List, Tuple

import pandas as pd
from sparkai.core.messages import ChatMessage

from iflytech_assistant.client import GenerationResult
from iflytech_assistant.client import instruct_with_profile as instruct

from .dataclasses import FiveSuggestions, PreferenceTag
from .markdown_prompt_parser import PromptTemplate

__all__ = ["generate", "generate_response", "refine"]

templates_dir = os.path.join(os.path.dirname(__file__), "templates")
template_file = os.path.join(templates_dir, "heqc.md")
prompt_template = PromptTemplate.from_markdown_text(open(template_file).read())

SYSTEM_PROMPT = prompt_template.system_prompt_template.raw_str
REPLY_USER_PROMPT = prompt_template.reply_user_prompt_template.raw_str
POLISH_USER_PROMPT = prompt_template.polish_user_prompt_template.raw_str
REFINE_PROMPT = prompt_template.refine_user_prompt_template.raw_str


def build_prompt(
    user_input: str, target: str, tags: List[PreferenceTag], mode: str
) -> Tuple[str, str]:
    tone_prompt = "\n".join([f"- {tag.tag}: {tag.tone}" for tag in tags if tag.tone])
    style_prompt = "\n".join([f"- {tag.tag}: {tag.style}" for tag in tags if tag.style])
    personality_prompt = "\n".join(
        [f"- {tag.tag}: {tag.personality}" for tag in tags if tag.personality]
    )

    system_prompt = SYSTEM_PROMPT.format(
        tone=tone_prompt, style=style_prompt, personality=personality_prompt
    )

    if mode == "reply":
        user_prompt = REPLY_USER_PROMPT.format(
            target=target,
            user_input=user_input,
            preferenced_style="、".join([tag.tag for tag in tags]),
        )
    else:
        user_prompt = POLISH_USER_PROMPT.format(
            target=target,
            user_input=user_input,
            preferenced_style="、".join([tag.tag for tag in tags]),
        )
    return system_prompt, user_prompt


def generate(user_input, target: str, tags: List[PreferenceTag], mode: str):
    system_prompt, user_prompt = build_prompt(user_input, target, tags, mode)
    llm_response = instruct(user_prompt, system_prompt, prefix="")

    suggestions = FiveSuggestions.from_generation_result(llm_response)

    return system_prompt, user_prompt, llm_response, suggestions


def refine(
    user_input: str,
    system_prompt: str,
    refine_words: str,
    llm_response: GenerationResult,
):
    history = [
        ChatMessage(role="user", content=user_input),
        ChatMessage(role="assistant", content=llm_response.llm_response),
    ]

    refine_prompt = REFINE_PROMPT.format(refine_words=refine_words)
    try:
        llm_refine_response = instruct(
            refine_prompt, system_prompt, prefix="", historys=history
        )
    except Exception:
        llm_refine_response = GenerationResult(llm_response="")

    refine_suggestions = FiveSuggestions.from_generation_result(llm_refine_response)

    return refine_suggestions, refine_prompt


def generate_response(
    user_input: str,
    target: str,
    tags: List[PreferenceTag],
    refine_words: str,
    mode: str,
) -> Tuple[FiveSuggestions, FiveSuggestions]:
    system_prompt, user_prompt, llm_response, suggestions = generate(
        user_input, target, tags, mode
    )

    refine_suggestions, refine_prompt = refine(
        user_input, system_prompt, refine_words, llm_response
    )
    return (
        suggestions,
        refine_suggestions,
        system_prompt,
        user_prompt,
        refine_prompt,
    )


def load_tags_from_excel(
    file_path: str,
) -> Tuple[Dict[str, List[PreferenceTag]], Dict[str, List[str]]]:
    # read all sheet names
    xls = pd.ExcelFile(file_path)
    sheets = xls.sheet_names

    assert "修改提示词" in sheets, "Sheet '修改提示词' not found"

    refine_words_sheet = xls.parse("修改提示词")
    # remove first row
    refine_words_sheet = refine_words_sheet.iloc[1:]
    # remove Unnamed columns
    refine_words_sheet = refine_words_sheet.loc[
        :, ~refine_words_sheet.columns.str.contains("^Unnamed")
    ]
    # only keep columns with "高情商-{target}"
    refine_words_sheet = refine_words_sheet.loc[
        :, refine_words_sheet.columns.str.contains(r"高情商-.*")
    ]
    # column names 高情商-{target} -> target
    refine_words_sheet.columns = refine_words_sheet.columns.str.extract(r"高情商-(.*)")[
        0
    ]

    refine_words = {
        col: refine_words_sheet[col].dropna().tolist()
        for col in refine_words_sheet.columns
    }

    tags = {}

    for sheet in sheets:
        if not re.match(r"生成偏好-.*", sheet):
            continue
        target = sheet.split("-")[1]
        tags[target] = xls.parse(sheet)

    tag_objs = {}

    for target, tag_sheet in tags.items():
        tag_objs[target] = []
        for _, row in tag_sheet.iterrows():
            # TODO spell error
            personality = row["Personility"]
            tone = row["Tone"]
            style = row["Style"]

            tag_objs[target].append(
                PreferenceTag(
                    tag=row.iloc[0], tone=tone, style=style, personality=personality
                )
            )

    return tag_objs, refine_words
